import torch
from torch import nn
from torch_geometric import nn as nng
import sys
import os
cur_path = os.path.split(__file__)[0]
sys.path.append(cur_path)
sys.path.append(os.path.join(cur_path, '..'))
from UNet.unet_model import UNet
import torch.nn.functional as F



def build_mlp(
    in_size, hidden_size, out_size, drop_out=True, lay_norm=True, dropout_prob=0.2
):
    if drop_out:
        module = nn.Sequential(
            nn.Linear(in_size, hidden_size),
            nn.Dropout(p=dropout_prob),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.Dropout(p=dropout_prob),
            nn.SiLU(),
            nn.Linear(hidden_size, out_size),
        )
    else:
        module = nn.Sequential(
            nn.Linear(in_size, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.SiLU(),
            nn.Linear(hidden_size, out_size),
        )
    if lay_norm:
        return nn.Sequential(module, nn.LayerNorm(normalized_shape=out_size))
    return module



class Model(nn.Module):
    def __init__(self, params, bounds) -> None:
        super().__init__()
        # self.encoder = build_mlp(1, 32, 8,lay_norm=True,drop_out=False)
        self.unet = UNet(1, 16)
        self.decoder = build_mlp(16, 64, 1,lay_norm=False,drop_out=False)

        self.params = params
        self.bounds = torch.from_numpy(bounds).to(torch.float32)
        # self.resolution = resolution
        # self.feature_mean_std = feature_mean_std
        # self.label_mean_std = label_mean_std

    
    def forward(self, batch_data):
        # B C D H W
        batch_size = batch_data['node|pos'].shape[0]

        # z = self.encoder(torch.permute(batch_data['voxel|sdf'], [0, 2, 3, 4, 1])) # B D H W C
        # z = torch.permute(batch_data['voxel|sdf'], [0, 2, 3, 4, 1])
        # latent_voxel = self.unet(torch.permute(z, [0, 4, 1, 2, 3]))
        z = batch_data['voxel|sdf']
        latent_voxel = self.unet(z)
        
        z_out = F.grid_sample(latent_voxel, batch_data['node|pos']) 
        z_out = z_out.reshape(batch_size, batch_data['node|pos'].shape[3], -1)
        output = self.decoder(z_out)
        # output = self.shitTest(batch_data['node|pos'].reshape(batch_size, batch_data['node|pos'].shape[3], -1))
        return output